from langchain_community.llms import Ollama
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from business.utils.vector_store import VectorStore
from business.ollama.llm_check import LLMCheck
from config.index import Config
from typing import AsyncGenerator, Dict, Any
import json
import os
import asyncio

class OllamaChatDoc:
    def __init__(self):
        self.cfg = Config()
        self.host = self.cfg.get('ollama.host')
        self.model = self.cfg.get('ollama.model')
        
        # 初始化 deepseek 模型，启用流式输出
        self.llm = Ollama(
            base_url=self.host,
            model=self.model,
            callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
            temperature=self.cfg.get('ollama.temperature'),
        )
        
        # 初始化LLM检查器（使用 phi3）
        self.llm_check = LLMCheck()
        
        # 向量存储相关
        self.vector_store = None
        self._vector_store_initialized = False

    def ensure_vector_store(self) -> bool:
        """确保向量存储已初始化（仅在需要时初始化）"""
        if self._vector_store_initialized:
            return self.vector_store is not None
            
        try:
            vector_store_path = os.path.abspath("data/vector")
            if not os.path.exists(vector_store_path):
                print("警告：向量库不存在，请先运行 rag_init.py 初始化向量库")
                self._vector_store_initialized = True
                return False
                
            print("正在加载向量库...")
            vs = VectorStore()
            self.vector_store = vs.load_vector_store(vector_store_path)
            self._vector_store_initialized = True
            
            if self.vector_store:
                print("向量库加载成功！")
                return True
            else:
                print("警告：向量库加载失败")
                return False
        except Exception as e:
            print(f"加载向量库时发生错误: {e}")
            self._vector_store_initialized = True
            return False

    def get_relevant_context(self, query: str, k: int = 5) -> str:
        """从向量库中检索相关上下文"""
        if not self.ensure_vector_store():
            return ""
            
        try:
            # 预处理查询词
            clean_query = query.replace('-', ' ').replace('_', ' ')  # 替换特殊字符
            terms = clean_query.split()  # 分词
            
            # 构建搜索关键词
            search_terms = [clean_query]  # 原始查询（已清理）
            
            # 如果是多个词，也尝试单独搜索每个词
            if len(terms) > 1:
                search_terms.extend(terms)
            
            # 添加其他搜索变体
            for term in list(search_terms):  # 使用 list() 创建副本以便在循环中添加
                if term.strip():
                    search_terms.extend([
                        f"### {term}",  # 作为API标题搜索
                        f"{term} API",   # API相关
                        f"{term} 使用方法",  # 使用方法
                        f"{term} 示例"   # 示例代码
                    ])
            
            # 去除空字符串并去重
            search_terms = list(set(term.strip() for term in search_terms if term.strip()))
            
            all_results = []
            seen_contents = set()  # 用于去重
            
            # 对每个搜索词进行检索
            for term in search_terms:
                try:
                    results = self.vector_store.similarity_search(term, k=2)
                    for doc in results:
                        content = doc.page_content
                        # 去重并保持相关性
                        if content not in seen_contents:
                            seen_contents.add(content)
                            all_results.append(doc)
                except Exception as e:
                    print(f"搜索词 '{term}' 检索失败: {e}")
                    continue
            
            if not all_results:
                return ""
                
            # 将检索结果组合成上下文
            context_parts = []
            for i, doc in enumerate(all_results, 1):
                source = doc.metadata.get('source', '未知来源') if hasattr(doc, 'metadata') else '未知来源'
                context_parts.append(f"文档 {i} (来源: {source}):\n{doc.page_content}\n")
            
            return "\n".join(context_parts)
        except Exception as e:
            print(f"检索相关内容时发生错误: {e}")
            return ""

    def _build_prompt(self, query: str, context: str = "", use_context: bool = True) -> str:
        """构建提示词"""
        if context and use_context:
            return f"""你是一个专业的Uniapp开发助手。请基于以下知识库回答用户的问题。

            参考文档:
            {context}

            用户问题: {query}
            
            规则：
            1. 请使用markdown格式回答，确保代码示例使用正确的代码块标记。
            2. 如果文档信息不足，请明确指出缺失的信息。
            3. 不用重复显示用户的问题是什么，只需要直接回答问题就行。"""
        else:
            return f"""你是一个专业的AI助手。请回答用户的问题。
            用户问题: {query}
            
            规则：
            1. 不用重复显示用户的问题是什么，只需要直接回答问题就行。"""

    async def stream_chat(self, query: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        流式聊天响应
        """
        try:
            # 检查是否需要知识库检索（使用 phi3）
            needs_retrieval = self.llm_check.needs_retrieval(query)
            print(f"是否需要知识库检索: {needs_retrieval}")
            # 获取相关上下文（如果需要）
            context = self.get_relevant_context(query) if needs_retrieval else ""
            print(f"相关上下文: {context}")
            # 构建提示词
            prompt = self._build_prompt(query, context, needs_retrieval)
            print(f"提示词: {prompt}")
            # 使用 deepseek 模型生成回答
            for chunk in self.llm.stream(prompt):
                yield {
                    "event": "message",
                    "data": json.dumps({
                        "content": chunk,
                        "type": "content"
                    }, ensure_ascii=False)
                }
                await asyncio.sleep(0.1)

        except Exception as e:
            yield {
                "event": "error",
                "data": json.dumps({
                    "error": str(e),
                    "type": "error"
                }, ensure_ascii=False)
            }

    def chat(self, query: str) -> str:
        """
        普通聊天（非流式），返回完整响应
        """
        try:
            # 检查是否需要知识库检索（使用 phi3）
            needs_retrieval = self.llm_check.needs_retrieval(query)
            
            # 获取相关上下文（如果需要）
            context = self.get_relevant_context(query) if needs_retrieval else ""
            
            # 如果找到了上下文且需要使用，打印出来
            if context and needs_retrieval:
                print("\n找到的相关内容:")
                print("=" * 40)
                print(context)
                print("=" * 40)
            
            # 构建提示词
            prompt = self._build_prompt(query, context, needs_retrieval)
            
            print("\nAI助手正在思考...")
            # 使用 deepseek 模型生成回答
            response = self.llm.invoke(prompt)
            return response
            
        except Exception as e:
            print(f"对话过程中发生错误: {e}")
            return str(e)
